import time
import requests
import pandas as pd
from QFinanceGridModel.base import generate_ut_token, Url, Headers
from tqdm import tqdm  # 用于显示进度条
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def get_one_page_flow_direction_ranking(sector_code, page, page_size=100):
    # 确定fs参数和url_code
    if sector_code in {"industry", "2"}:
        url = Url(type="2").base_url
        params = Url(type="2").params
        headers = Headers(type="2").headers
    elif sector_code in {"concept", "3"}:
        url = Url(type="3").base_url
        params = Url(type="3").params
        headers = Headers(type="3").headers
    else:
        url = Url(type="4", secid=sector_code).base_url
        params = Url(type="4", secid=sector_code).params
        headers = Headers(type="4", secid=sector_code).headers

    try:
        response = requests.get(
            url=str(url),
            params=params,
            headers=headers,
            timeout=30
        )
        # print(response.url)
        response.raise_for_status()
        data = response.json()
        # print(data)
        # 检查数据有效性
        if 'data' not in data or 'diff' not in data['data']:
            logger.warning(f"第 {page} 页未获取到有效数据")
            return 0, None

        total_count = data['data']['total']
        page_count = (total_count + page_size - 1) // page_size

        # 处理数据
        records = []
        for idx, item in enumerate(data['data']['diff'], start=1):
            record = {
                "代码": str(item.get("f13", ""))+"."+item.get("f12", ""),  # f13 =1 沪，=0 深，90表示板块/行业/概念
                "名称": item.get("f14", ""),
                "最新价": item.get("f2", 0),
                "今日涨幅(%)": item.get("f3", 0),
                "主力净流入最大个股代码": item.get("f205", ""),
                "主力净流入最大个股名称": item.get("f204", ""),
                "主力净流入(元)": item.get("f62", 0),
                "主力净占比(%)": item.get("f184", 0),
                "超大单净流入(元)": item.get("f66", 0),
                "超大单净占比(%)": item.get("f69", 0),
                "大单净流入(元)": item.get("f72", 0),
                "大单净占比(%)": item.get("f75", 0),
                "更新时间": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(item.get("f124", 0)))
            }
            records.append(record)

        df = pd.DataFrame(records)
        return page_count, df

    except requests.exceptions.RequestException as e:
        logger.error(f"请求失败: {str(e)}")
    except (ValueError, KeyError) as e:
        logger.error(f"数据处理错误: {str(e)}")
    except Exception as e:
        logger.error(f"未知错误: {str(e)}")

    return 0, None


def get_all_flow_direction_ranking(sector_code):
    """
    获取所有分页数据
    参数:
        sector_code: 行业/概念代码或类型标识
    返回:
        DataFrame: 包含所有分页数据的DataFrame
    """
    # 获取第一页数据
    page_count, first_page_df = get_one_page_flow_direction_ranking(sector_code, 1)

    if first_page_df is None:
        logger.error("获取第一页数据失败")
        return pd.DataFrame()

    if page_count <= 1:
        return first_page_df

    # 收集所有分页数据
    all_dfs = [first_page_df]

    # 使用进度条显示获取进度
    for page in tqdm(range(2, page_count + 1), desc="获取数据", unit="页"):
        _, page_df = get_one_page_flow_direction_ranking(sector_code, page)
        if page_df is not None and not page_df.empty:
            all_dfs.append(page_df)
        time.sleep(0.5)  # 添加请求间隔防止被封

    # 合并所有数据
    return pd.concat(all_dfs, ignore_index=True)


if __name__ == "__main__":
    # 设置显示选项
    pd.set_option('display.width', 1000)
    pd.set_option('display.max_columns', None)
    pd.set_option('display.float_format', lambda x: f"{x:,.2f}")

    # 执行数据获取
    df = get_all_flow_direction_ranking("BK0947")  # BK0547 BK1106

    if not df.empty:
        print("\n获取结果:")
        print(df.head(20))
        print(f"\n总共获取 {len(df)} 条记录")
    else:
        print("未获取到有效数据")
